package utils
import chisel3._
import chisel3.util._

class SRAMWriteIO[T <: Data](private val gen: T, val set: Int /*, val way: Int = 1*/) extends Bundle {
  val setIdx = Input(UInt(log2Up(set).W))
  val data = Input(gen)
  val wen = Input(Bool())
  //val waymask = if (way > 1) Some(Input(UInt(way.W))) else None
}
class SRAMReadIO[T <: Data](private val gen: T, val set: Int /*, val way: Int = 1*/) extends Bundle {
  val setIdx = Input(UInt(log2Up(set).W))
  val data = Output(gen)
  val ren = Input(Bool())
}

/**
 * SRAM Template
 */
class SRAM[T <: Data](
private val gen: T,
set: Int,
way: Int = 1,
holdRead: Boolean = false,
singlePort: Boolean = false
) extends Module{
  val io = IO(new Bundle() {
    val w = new SRAMWriteIO(gen, set)
    val r = new SRAMReadIO(gen, set)
  })

  val wordType = UInt(gen.getWidth.W)
  // sram array
  val ram = SyncReadMem(set, wordType)

  val ren =  io.r.ren
  val wen = io.w.wen
  val realRen = (if (singlePort) ren && !wen else ren)

  val setIdx = io.w.setIdx
  val wdataword = io.w.data.asUInt

  when (wen) {
    ram.write(setIdx, wdataword)
  }

  val rdata = if(holdRead) ReadAndHold(ram, io.r.setIdx, realRen)
  else ram.read(io.r.setIdx, realRen)

  io.r.data := rdata.asTypeOf(gen)

  when (reset.asBool()) {
    for (i <- 0 until set) {
      ram.write(i.U, 0.U)
    }
  }
}
